#lang racket

(require "util/utils.scm")
(require "util/match-assoc.scm")
;;(require "util/collector.scm")
(provide pass-labels)

;; this is just a slighlty modified version of labels
;; used for c-gen instead of vectorcode

(define (make-label-collector)
  (let ((labels '()))
    (lambda (label)
      (cond ((assoc label labels) => cdr)
            (else (let ((blabel (box #f)))
                    (set! labels (cons (cons label blabel) labels))
                    blabel))))))

(define (pass-labels program)
  (match-assoc program (entry-point globals-size code)
    (let* ((genlabel! (make-label-collector))
           (entry-point (genlabel! entry-point))
           (code-boxed (box-labels code 0 (make-collector) genlabel!))
           (entry-point (list (unbox entry-point)))
           (code-unboxed (debox-labels code-boxed (make-collector))))
      `((entry-point . ,entry-point)
        (globals-size . ,globals-size)
	(code . ,code-unboxed)))))

(define (box-labels exps pos code genlabel!)
  (if (null? exps)
      (collector-get code)
      (let ((pos* (box-exp (car exps) pos code genlabel!)))
        (box-labels (cdr exps) pos* code genlabel!))))

(define (box-exp exp pos code genlabel!)
  (match exp
    (`(label ,nm)
     (set-box! (genlabel! nm) pos)
     (collect-object! code `(label ,pos)) ;; CHANGED
     (+ pos 1) ;; CHANGED
     )
    (`(push! ,x)
     (collect-object! code `(push! ,(box-exp-simple x genlabel!)))
     (+ pos 1))
    (`(branch ,test ,label)
     (collect-object! code `(branch ,(box-exp-simple test genlabel!) ,(genlabel! label)))
     (+ pos 1))
    (`(move! ,to ,from)
     (collect-object! code `(move! ,to ,(box-exp-simple from genlabel!)))
     (+ pos 1))
    (`(jump ,label)
     (collect-object! code `(jump ,(genlabel! label)))
     (+ pos 1))
    (`(closure-set! ,c ,n ,v)
     (collect-object! code `(closure-set! ,(box-exp-simple c genlabel!)
                                          ,n
                                          ,(box-exp-simple v genlabel!)))
     (+ pos 1))
    (`(call)
     (collect-object! code `(call))
     (+ pos 1))
    (`(shrink-stack ,n)
     (collect-object! code exp)
     (+ pos 1))
    (`(extend-stack ,n)
     (collect-object! code exp)
     (+ pos 1))
    (`(stack-fold ,w ,d)
     (collect-object! code exp)
     (+ pos 1))
    (`(return)
     (collect-object! code exp)
     (+ pos 1))
    (`(exit)
     (collect-object! code exp)
     (+ pos 1))
    (else (error 'box-exp "?" exp))))

(define (box-exp-simple exp genlabel!)
  (match exp
    (`(quote ,x) exp)
    (`(label ,nm) `(label ,(genlabel! nm)))
    (`(env-register) exp)
    (`(return-register) exp)
    (`(atomic ,v) exp)
    (`(variable ,ty ,loc) exp)
    (`(builtin ,nm . ,args)
     `(builtin ,nm . ,(mapply box-exp-simple args genlabel!)))
    (`(closure ,env (label ,nm))
     `(closure ,env (label ,(genlabel! nm))))
    (else (error 'box-exp-simple "?" exp))))

(define (debox-labels p code)
  (if (null? p)
      (collector-get code)
      (let ((pos* (debox-exp (car p) code)))
        (debox-labels (cdr p) code))))

(define (debox-exp exp code)
  (match exp
    (`(call)
     (collect-object! code `(call)))
    (`(push! ,x)
     (collect-object! code `(push! ,(debox-exp-simple x))))
    (`(move! ,to ,from)
     (collect-object! code `(move! ,to ,(debox-exp-simple from))))
    (`(branch ,test ,blabel)
     (collect-object! code `(branch ,(debox-exp-simple test) ,(unbox blabel))))
    (`(jump ,blabel)
     (collect-object! code `(jump ,(unbox blabel))))
    (else (collect-object! code exp))))

(define (debox-exp-simple exp)
  (match exp
    (`(label ,blabel) `(label ,(unbox blabel)))
    (`(quote ,x) exp)
    (`(env-register) exp)
    (`(return-register) exp)
    (`(atomic ,v) exp)
    (`(variable ,ty ,loc) exp)
    (`(builtin ,nm . ,args)
     `(builtin ,nm . ,(mapply debox-exp-simple args)))
    (`(closure ,env (label ,blabel))
     `(closure ,env (label ,(unbox blabel))))
    (else (error 'debox-exp-simple "?" exp))))
